# -*- coding:utf-8 -*-

# @Time    : 2023/5/16 02:21
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : generate_bot_dialogue.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

import sys
import os

current_script_path = os.path.abspath(__file__)
work_path = os.path.normpath(os.path.join(current_script_path, '../../'))
sys.path.append(work_path)
sys.path.append(work_path + '/internal_server')

import pandas as pd
import asyncio
from bot.insurance_planner_gpt.planner import PlannerChat
import uuid
from data_generate import utils
import traceback
import random
import os
from util_tool import utils
from bot.insurance_planner_gpt.make_local_data.user import User
from bot.insurance_planner_gpt.utils.conversation_process_util import ConversationProcessUtil
from bot.insurance_planner_gpt.context_local_util import context_data


# gpt4模拟用户生成bot对话
async def mock_bot_dialogue(path, taget_num=6, round_num=15, start_context=[]):
    train_data = []
    if os.path.exists(path):
        train_data = utils.jload(path)

    out_writer = open("dialogue_tmp.txt", 'a+')
    for i in range(taget_num):
        try:
            session_id = str(uuid.uuid1()).replace('-', '')
            context_data.set({"session_id": session_id, "message_id": 'message_id'})
            # local_data.session_id = session_id

            context = start_context.copy()
            instruct_dict_new = {}
            instruct_dict_new["id"] = session_id
            instruct_dict_new["messages"] = []
            instruct_dict_new["messages"].append({
                "role": "user",
                "content": ""
            })
            for i in range(round_num):
                message_id = str(uuid.uuid1()).replace('-', '')
                context_data.set({"session_id": session_id, "message_id": message_id})
                # local_data.message_id = message_id
                history = ConversationProcessUtil.format_conversation_history(context, assistant_role_name='保险助理')
                user = User(history)
                # 如果是用户第一句，那么在这里选开头
                # if '用户' not in history:
                #     query = '嗨，犀心小助，最近我挺关心健康方面的问题，想给自己买个保险。我30岁，软件工程师这边的，有糖尿病和高血脂，我能买些什么医疗险？哦，对了，我没有社保。'
                # else:
                #     query = await user.achat_auto_llm(type="gpt")
                query = await user.achat_auto_llm(type="gpt")
                conversation_dict = {}
                conversation_dict["role"] = "user"
                conversation_dict["content"] = query.replace("<end>", "")
                instruct_dict_new["messages"].append(conversation_dict)
                context.append(conversation_dict)

                bot = PlannerChat(None,
                                  planner_role="保险助理",
                                  company_name="慧择保险网",
                                  company_business="仅为用户提供全面的保险规划服务，帮助用户解决任何保险规划或保险产品等相关的问题，以解决用户问题为目标，不以销售产品为导向。",
                                  company_values="真诚、温暖、专业，做用户的真朋友，尊重每个用户个体的差异，用心为用户提供最优质的服务。"
                                  )

                res, other_messages = await bot.async_reply(context, session_id, None)
                conversation_dict = {}
                conversation_dict["role"] = "assistant"
                conversation_dict["content"] = res
                instruct_dict_new["messages"].append(conversation_dict)
                context.append(conversation_dict)

                if "<end>" in query:
                    break
            train_data.append(instruct_dict_new)
            utils.jdump(train_data, path)
            out_writer.write(str(instruct_dict_new) + "\n")
            out_writer.flush()
        except Exception as e:
            traceback.print_exc()
            # 打印堆栈
            print(e)

    out_writer.close()


def add_empty_row(group):
    return pd.concat([group, pd.DataFrame([[''] * len(group.columns)], columns=group.columns)],
                     ignore_index=True)


def conversation2csv(file_path):
    datas = utils.jload(file_path)
    df = pd.DataFrame(columns=['角色', '内容'])

    for data in datas:
        messages = data['messages']
        for text in messages:
            df = df.append(
                {"角色": text['role'],
                 "内容": text['content']},
                ignore_index=True)
        df = df.reset_index(drop=True).append(
            {"角色": "",
             "内容": ""},
            ignore_index=True)
    df.to_csv("gen_conversation.csv", index=False, encoding='utf-8_sig')


def read_agent_info(file_path):
    datas = utils.jload(file_path)
    resp = {}
    for data in datas:
        _id = data["id"]
        _time = data["time"]
        _task = data["task"]
        _model_name = data["model_name"]
        _conversations = data["conversations"]
        for conv in _conversations:
            if conv["from"] == 'gpt':
                value = conv["value"]
        if _id not in resp:
            resp[_id] = {}
        if _task not in resp[_id]:
            resp[_id][_task] = {}
        resp[_id][_task]["model_name"] = _model_name
        resp[_id][_task]["time"] = _time
        resp[_id][_task]["value"] = value

    return resp


def json2xlsx(json_path, excel_path):
    data = utils.jload(json_path)
    # 准备一个空列表来存储解析后的消息
    messages_list = []

    # 遍历每个会话及其消息
    for conversation in data:
        conversation_id = conversation['id']
        for message in conversation['messages']:
            # 创建一个包含id, role, content的字典，并将其添加到列表中
            message_dict = {
                'id': conversation_id,
                'role': message['role'],
                'content': message['content']
            }
            messages_list.append(message_dict)

        # 将列表转换为DataFrame
    df = pd.DataFrame(messages_list)

    # 将DataFrame写入Excel文件，不包含索引列
    df.to_excel(excel_path, index=False)


def conversation2xlsx(dialogue_path, agent_path):
    agent_res = read_agent_info(agent_path)
    datas = utils.jload(dialogue_path)
    ids = []
    roles = []
    contents = []
    model_names = []
    times = []
    # df = pd.DataFrame(columns=['对话id', '角色', '对话内容', '对话时间', '模型名称'])
    for data in datas:
        _id = data["id"]
        _messages = data["messages"]
        for text in _messages:
            role = text['role']
            content = text['content']
            if "message_id" in text:
                key = _id + ":" + text["message_id"]
                values = agent_res[key]
                for k, v in values.items():
                    model_name = v["model_name"]
                    time = v["time"]
                    value = v["value"]
                    title = ""
                    if k == "QuestioningDisputeResolution":
                        title = "用户问题"
                    if k == "UserInfoExtract":
                        title = "用户信息"
                    if k == "ProductCausalSolution":
                        title = "因果推断"
                    if k == "PlanExplain":
                        title = "方案对比"
                    if k == "PlanAssume":
                        title = "因果图"
                    if title != "":
                        ids.append(_id)
                        roles.append(role)
                        contents.append({"title": title, "content": value})
                        model_names.append(model_name)
                        times.append(time)

                ids.append(_id)
                roles.append(role)
                contents.append(content)
                model_names.append("")
                times.append(time)
            else:
                ids.append(_id)
                roles.append(role)
                contents.append(content)
                model_names.append("")
                times.append("")

        ids.append("")
        roles.append("")
        contents.append("")
        model_names.append("")
        times.append("")

    df = pd.DataFrame({"对话id": ids, "角色": roles, "对话内容": contents, "对话时间": times, "模型名称": model_names})
    df.to_excel("conversation_1115_1.xlsx", index=False)


def filter_train_data(path, target_path):
    count = 0
    for file in os.listdir(path):
        if not file.endswith('.json'):
            continue
        if "conversation" in file:
            data = utils.jload(path + file)
            data = data[0:1000]
        else:
            data = utils.jload(path + file)
        result_data = []

        product_names = ["暖宝保", "小医仙", "长相安", "医享无忧", "众民保", "卓越馨选", "臻爱无忧", "MSH", "成长优享", "金医保",
                         "达尔文", "小红花", "i无忧", "超级玛丽", "守卫者", "康顺人生", "康乐一生", "小青龙", "小淘气", "麦兜兜",
                         "大麦", "臻爱2023", "小蜜蜂", "小神童", "孝心安", "小团圆", "金禧", "富德生命", "慧选", "智慧领航",
                         "臻享一生", "金满意足", "鑫相守", "中意一生", "平安去旅行", "一日游", "众行天下", "天平乐"]

        for record in data:
            new_record = {}
            conversations = record['conversations']
            if "model_name" in record:
                model_name = record['model_name']
            else:
                model_name = "gpt-4-conversations"
            if "gpt-4" not in model_name:
                continue
            if len(conversations) == 2:
                prompt = conversations[0]['value']
                answer = conversations[1]['value']
                if "的json格式是：" in prompt and "2024-02-01.json" in file:
                    continue
                if "的json格式是：" in prompt and "2024-01" in file:
                    continue
                if "用户:养老还需要什么" in prompt:
                    continue

                # '补充下一步和用户沟通策略的因果图：用户:你好 -> 了解用户保险需求（需用户提供）'
                # if "补充下一步和用户沟通策略的因果链：" in prompt and ("无" != answer):
                #     results = str(answer).split("\n")
                #     for result in results:
                #         if (("（需讲解）" not in result) and ("（需用户提供）" not in result) and ("（需计算）" not in result) and ("（潜在需求）" not in result)  and ("conversation" not in file)):
                #             print(answer)

                # 保险规划_如何选择医疗险->中端医疗（保险产品类型）->慧择小医仙2号医疗险-计划一（保险产品名称）（保险类型：医疗险，保险产品类型：小额医疗险）
                if ("无" != answer):

                    results = str(answer).split("\n")
                    for result in results:
                        if (("中端" in result) and ("小医仙" in result)):
                            print(answer)
                # if ("（阶段" in answer):
                #     print(answer)

                # if ("所有需要识别的内容" in prompt):
                #     no_example_prompt = """示例（任务忽略）：\n---\n用户:我今年本科毕业后第三年\n每年收入240000元 \n每月支出1000元  //当前年支出 1000*12=12000，不要返回计算过程，只返回12000\n没有存款\n父母年龄分别是55和53\n每年保费预算5000\n两个女孩，8岁和3岁\n我想看看小蜜蜂意外险\n---\n输出的json结果： {\"本人年龄\": 26, \"当前年收入\": 240000, \"当前年支出\": 12000, \"当前存款\": 0,\n\"父母年龄\":[55, 53], \"年保费预算\": 5000,\n\"孩子个数\":2, \"孩子年龄\":[8,3], \"孩子性别\":[女,女], \"孩子人生阶段\": [\"大龄儿童\", \"低龄儿童\"], \"当前关注保险产品\":\"小蜜蜂意外险\"}\n示例结束;"""
                #     is_in = False
                #     for product_name in product_names:
                #         if  (product_name in str(prompt).replace(no_example_prompt, "").replace('当前关注保险产品"是指保险产品的名称，比如"暖宝保',"").replace('当前关注保险产品"是指保险产品的名字，比如"暖宝保',"").replace("'当前关注保险产品'是指保险产品的名称，比如'暖宝保'", "")) and ("当前关注保险产品" not in answer):
                #
                #             if "2024-01-24" in file:
                #                 is_in = True
                #
                #                 if count > 9:
                #                     print(file)
                #
                #                     print(answer)
                #     if is_in:
                #         count = count + 1

            new_record['id'] = str(uuid.uuid1())
            new_record['model_name'] = model_name
            new_record["file_name"] = str(file)
            new_record['conversations'] = conversations
            result_data.append(new_record)

        if result_data:
            utils.jdump(result_data, target_path + file)


def merge_train_data(path, target_path):
    datas = []
    for file in os.listdir(path):
        if not file.endswith('.json'):
            continue
        data = utils.jload(path + file)

        for record in data:
            new_record = {}
            conversations = record['conversations']
            if len(conversations) == 2:
                prompt = conversations[0]['value']
                answer = conversations[1]['value']
                if "transfer_flag" in prompt:
                    continue
                if "推荐用户接下来沟通内容" in prompt:
                    continue
                if "请输出你总结的内容" in prompt:
                    continue

            new_record['id'] = str(uuid.uuid1())
            new_record["file_name"] = str(file)
            new_record['conversations'] = conversations
            datas.append(new_record)
    # 打乱datas
    random.shuffle(datas)
    print(len(datas))
    utils.jdump(datas, target_path + "planner_train_data.json")


def csv2vicuna(file_paths, target_path):
    train_datas = []
    for file_path in file_paths:
        df = pd.read_csv(file_path)
        # df = df.dropna()
        for index, row in df.iterrows():
            train_data = {}
            session_id = str(uuid.uuid1())
            conversations = []
            conversation = {}
            if pd.isna(row['prompt']) or pd.isna(row['修改后结果']):
                print(row)
                continue

            conversation['from'] = 'human'
            conversation['value'] = row['prompt']
            conversations.append(conversation)
            conversation = {}
            conversation['from'] = 'gpt'
            conversation['value'] = row['修改后结果']
            conversations.append(conversation)
            train_data['id'] = session_id
            train_data['conversations'] = conversations
            train_datas.append(train_data)
    utils.jdump(train_datas, target_path)


def knowledge_csv2vicuna(file_paths, target_path):
    train_datas = []
    for file_path in file_paths:
        df = pd.read_csv(file_path)
        # df = df.dropna()
        for index, row in df.iterrows():
            train_data = {}
            session_id = str(uuid.uuid1())
            conversations = []
            conversation = {}
            if pd.isna(row['用户问题']) or pd.isna(row['答案']):
                print(row)
                continue

            conversation['from'] = 'human'
            conversation['value'] = row['用户问题']
            conversations.append(conversation)
            conversation = {}
            conversation['from'] = 'gpt'
            conversation['value'] = row['答案']
            conversations.append(conversation)
            train_data['id'] = session_id
            train_data['conversations'] = conversations
            train_datas.append(train_data)
    utils.jdump(train_datas, target_path)


if __name__ == '__main__':
    # nohup python ./data_generate/generate_planner_bot_dialogue_wechat.py >> qc_generate_planner_bot_dialogue_wechat-0512.log 2>&1 &
    import datetime

    # 按顺序跑个
    now_time = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M")
    file_path = "../data/planner/bot_dialogue/bot_dialogue_" + now_time + ".json"  # 文件夹路径
    excel_path = "../data/planner/bot_dialogue/bot_dialogue_" + now_time + ".xlsx"
    the_context =[
    [{'role': 'user', 'content': '我通过了你的联系人验证请求，现在我们可以开始聊天了'},
     {'role': 'assistant', 'content': '欢迎使用慧择保险服务，有任何保险相关的问题都可以联系慧择保险，慧择保险竭诚为您服务！'},
     {'role': 'assistant', 'content': '🔅您好呀～我是您的专属慧择小助理。\n您最近有什么想要关注和了解的保险么？\n回复【医疗险】【意外险】【重疾险】我给您做详细介绍，帮您分析适配度～'},
     ],
    [{'role': 'user', 'content': '我通过了你的联系人验证请求，现在我们可以开始聊天了'},
     {'role': 'assistant', 'content': '欢迎使用慧择保险服务，有任何保险相关的问题都可以联系慧择保险，慧择保险竭诚为您服务！'},
     {'role': 'assistant', 'content': '🔅您好呀～我是您的专属慧择小助理。\n您最近有什么想要关注和了解的保险么？\n回复【医疗险】【意外险】【重疾险】我给您做详细介绍，帮您分析适配度～'},
     ],
    [{'role': 'user', 'content': '我通过了你的联系人验证请求，现在我们可以开始聊天了'},
     {'role': 'assistant', 'content': '欢迎使用慧择保险服务，有任何保险相关的问题都可以联系慧择保险，慧择保险竭诚为您服务！'},
     {'role': 'assistant', 'content': '🔅您好呀～我是您的专属慧择小助理。\n您最近有什么想要关注和了解的保险么？\n回复【医疗险】【意外险】【重疾险】我给您做详细介绍，帮您分析适配度～'},
     {'role': 'user', 'content': '医疗险'},
     {'role': 'assistant', 'content': '好的，您是想给自己看呢还是给家里人看，年龄和我说下，我给您推荐适合的'},
     ],
    [{'role': 'user', 'content': '我通过了你的联系人验证请求，现在我们可以开始聊天了'},
     {'role': 'assistant', 'content': '欢迎使用慧择保险服务，有任何保险相关的问题都可以联系慧择保险，慧择保险竭诚为您服务！'},
     {'role': 'assistant', 'content': '🔅您好呀～我是您的专属慧择小助理。\n您最近有什么想要关注和了解的保险么？\n回复【医疗险】【意外险】【重疾险】我给您做详细介绍，帮您分析适配度～'},
     {'role': 'user', 'content': '意外险有什么'},
     {'role': 'assistant', 'content': '好的，您是想给自己看呢还是给家里人看，年龄和我说下，我给您推荐适合的'},
     ],
    [{'role': 'user', 'content': '我通过了你的联系人验证请求，现在我们可以开始聊天了'},
     {'role': 'assistant', 'content': '欢迎使用慧择保险服务，有任何保险相关的问题都可以联系慧择保险，慧择保险竭诚为您服务！'},
     {'role': 'assistant', 'content': '🔅您好呀～我是您的专属慧择小助理。\n您最近有什么想要关注和了解的保险么？\n回复【医疗险】【意外险】【重疾险】我给您做详细介绍，帮您分析适配度～'},
     {'role': 'user', 'content': '重疾'},
     {'role': 'assistant', 'content': '好的，您是想给自己看呢还是给家里人看，年龄和我说下，我给您推荐适合的'},
     ],
    [{'role': 'user', 'content': '我通过了你的联系人验证请求，现在我们可以开始聊天了'},
     {'role': 'assistant', 'content': '欢迎使用慧择保险服务，有任何保险相关的问题都可以联系慧择保险，慧择保险竭诚为您服务！'},
     {'role': 'assistant', 'content': '🔅您好呀～我是您的专属慧择小助理。\n您最近有什么想要关注和了解的保险么？\n回复【医疗险】【意外险】【重疾险】我给您做详细介绍，帮您分析适配度～'},
     {'role': 'user', 'content': '没时间'},
     {'role': 'assistant', 'content': '好的～我可以和您约一个您方便的时间么？比如明天上午的时间可以么？您可以险告诉我您想了解的，我提前准备发给您'},
     ],
    [{'role': 'user', 'content': '我通过了你的联系人验证请求，现在我们可以开始聊天了'},
     {'role': 'assistant', 'content': '欢迎使用慧择保险服务，有任何保险相关的问题都可以联系慧择保险，慧择保险竭诚为您服务！'},
     {'role': 'assistant', 'content': '🔅您好呀～我是您的专属慧择小助理。\n您最近有什么想要关注和了解的保险么？\n回复【医疗险】【意外险】【重疾险】我给您做详细介绍，帮您分析适配度～'},
     {'role': 'user', 'content': '医疗和重疾'},
     {'role': 'assistant', 'content': '好的，您是想给自己看呢还是给家里人看，年龄和我说下，我给您推荐适合的'},
     ]
    ]
    for the_context in the_context:
        asyncio.run(mock_bot_dialogue(file_path, taget_num=20, round_num=5,start_context=the_context))
    json2xlsx(file_path, excel_path)

